#!/usr/bin/env python3
"""
Author: Haoran Peng
Email: gavinsweden@gmail.com

An implementation of multi-agent path finding using conflict-based search
[Sharon et al., 2015]
"""

import multiprocessing as mp
from copy import deepcopy
from heapq import heappop, heappush
from itertools import combinations
from typing import Callable, Dict, List, Set, Tuple

import matplotlib.pyplot as plt
import numpy as np

# The low level planner for CBS is the Space-Time A* planner
# https://github.com/GavinPHR/Space-Time-AStar
from stastar.planner import Planner as STPlanner

from .agent import Agent
from .assigner import *
from .constraint_tree import CTNode
from .constraints import Constraints


class Planner:
    def __init__(
        self,
        grid_size: int,
        robot_radius: int,
        static_obstacles: List[Tuple[int, int]],
        num_agents: int,
    ):
        self.num_agents = num_agents
        self.robot_radius = robot_radius
        self.st_planner = STPlanner(grid_size, robot_radius, static_obstacles)

    def plan(
        self,
        starts: List[Tuple[int, int]],
        goals: List[Tuple[int, int]],
        assign_flag: int = 0,
        assign: Callable = min_cost,
        max_iter: int = 200,
        low_level_max_iter: int = 1000,
        max_process: int = 10,
        debug: bool = False,
    ) -> np.ndarray:
        self.low_level_max_iter = low_level_max_iter
        self.debug = debug
        # Do goal assignment
        self.agents = assign(starts, goals, assign_flag)
        self.agents_list = []

        for ag in self.agents:
            self.agents_list.append(ag.id)

        constraints = Constraints()

        # Compute path for each agent using low level planner
        solution = dict(
            (agent, self.calculate_path(agent, constraints, None))
            for agent in self.agents
        )

        open = []
        if all(len(path) != 0 for path in solution.values()):
            # Make root node
            node = CTNode(constraints, solution)
            # Min heap for quick extraction
            open.append(node)

        manager = mp.Manager()  # Manager for synchronization
        iter_ = 0
        while open and iter_ < max_iter:
            iter_ += 1
            print("Iteration", iter_, ":", max_iter)
            results = manager.list([])

            processes = []

            # Default to 10 processes maximum
            # for _ in range(max_process if len(open) > max_process else len(open)):
            #     p = mp.Process(target=self.search_node, args=[heappop(open), results])
            #     processes.append(p)
            #     p.start()

            # for p in processes:
            #     p.join()
            # results = []
            self.search_node(heappop(open), results)
            for result in results:
                if len(result) == 1:
                    if debug:
                        print(
                            "CBS_MAPF: Paths found after about {0} iterations".format(
                                iter_
                            )
                        )
                    return result[0]
                if result[0]:
                    heappush(open, result[0])
                if result[1]:
                    heappush(open, result[1])
        if debug:
            print("CBS-MAPF: Open set is empty, no paths found.")
        print("No paths found.")
        return np.array([])

    """
    Abstracted away the cbs search for multiprocessing.
    The parameters open and results MUST BE of type ListProxy to ensure synchronization.
    """

    def calculate_speed_fix(self, path, speed_fix):
        # 这里有对gird的处理，修改网格参数需要更新
        speed_fix = list(speed_fix)
        new_path = []
        for j in range(len(path)):
            for i in range(0, len(speed_fix)):
                x_point, y_point = path[j]
                new_path.append((x_point, y_point))
                speed_type = speed_fix[i][0]
                if speed_type == 0 or speed_type == 2:
                    continue
                start = speed_fix[i][1]
                end = speed_fix[i][3]
                if start <= x_point <= end:
                    if j < len(path) - 1:
                        next_x, next_y = path[j + 1]
                        if x_point != next_x:
                            for k in range(x_point + 5, next_x, 5):
                                new_path.append((k, y_point))
        return path if new_path == [] else np.array(new_path)

    def search_node(self, best: CTNode, results):
        (
            agent_i,
            agent_j,
            time_of_conflict,
            conflict_space,
        ) = self.validate_paths(self.agents, best)
        # If there is not conflict, validate_paths returns (None, None, -1)
        if agent_i is None:
            reformatted_solution = []
            for agent in self.agents:
                reformatted_solution.append(best.solution[agent])
            results.append((np.array(reformatted_solution),))
            return
        # Calculate new constraints
        agent_i_constraint = self.calculate_constraints(
            best,
            agent_i,
            agent_j,
            time_of_conflict,
            conflict_space,
            caclulate_type="i_j",
        )
        agent_j_constraint = self.calculate_constraints(
            best,
            agent_j,
            agent_i,
            time_of_conflict,
            conflict_space,
            caclulate_type="j_i",
        )

        # Calculate new paths

        fix_best_i = deepcopy(best.solution)
        # fix_best_i[agent_j] = self.calculate_speed_fix(
        #     best.solution[agent_j],
        #     agent_i_constraint.setdefault_speed(agent_j, set()),
        # )
        fix_best_i[agent_i] = self.calculate_path(
            agent_i,
            agent_i_constraint,
            self.calculate_goal_times(
                CTNode(agent_i_constraint, fix_best_i), agent_i, self.agents
            ),
        )

        fix_best_j = deepcopy(best.solution)
        # fix_best_j[agent_i] = self.calculate_speed_fix(
        #     best.solution[agent_i],
        #     agent_j_constraint.setdefault_speed(agent_i, set()),
        # )
        fix_best_j[agent_j] = self.calculate_path(
            agent_j,
            agent_j_constraint,
            self.calculate_goal_times(
                CTNode(agent_j_constraint, fix_best_j), agent_j, self.agents
            ),
        )
        # Replace old paths with new ones in solution

        node_i = None
        if all(len(path) != 0 for path in fix_best_i.values()):
            node_i = CTNode(agent_i_constraint, self.reformat(self.agents, fix_best_i))

        node_j = None
        if all(len(path) != 0 for path in fix_best_j.values()):
            node_j = CTNode(agent_j_constraint, self.reformat(self.agents, fix_best_j))
        results.append((node_i, node_j))

    """
    Pair of agent, point of conflict
    """

    def validate_paths(self, agents, node: CTNode):
        # Check collision pair-wise
        for agent_i, agent_j in combinations(agents, 2):
            # conflict_space = [conflict_type,ij,accelerate_obstacle,slowdown_obstacle]
            time_of_conflict, conflict_space = self.safe_distance(
                node.solution, agent_i, agent_j
            )
            # time_of_conflict=1 if there is not conflict
            if time_of_conflict == -1:
                continue
            return (
                agent_i,
                agent_j,
                time_of_conflict,
                conflict_space,
            )
        return None, None, -1, None

    def plot_conflict(self, ipath, jpath, agent_i, agent_j, idx, conflict_end_time):
        start_conflict_i = ipath[idx]
        start_conflict_j = jpath[idx]
        end_conflict_i = ipath[conflict_end_time]
        end_conflict_j = jpath[conflict_end_time]
        # 创建绘图
        plt.figure(figsize=(10, 6))

        # 绘制智能体i的路径
        plt.scatter(
            ipath[:, 0],
            ipath[:, 1],
            c="g",
            label=f"Agent {agent_i.id} Path",
        )
        plt.scatter(
            agent_i.start[0],
            agent_i.start[1],
            c="b",
            marker="o",
            label=f"Agent {agent_i.id} Start",
        )
        plt.scatter(
            agent_i.goal[0],
            agent_i.goal[1],
            c="b",
            marker="x",
            label=f"Agent {agent_i.id} Goal",
        )
        plt.scatter(
            start_conflict_i[0],
            start_conflict_i[1],
            c="b",
            marker="s",
            label=f"Agent {agent_i.id} Conflict Point",
        )
        plt.scatter(
            end_conflict_i[0],
            end_conflict_i[1],
            c="b",
            marker="p",
            label=f"Agent {agent_i.id} End Conflict",
        )

        # 绘制智能体j的路径
        plt.scatter(
            jpath[:, 0],
            jpath[:, 1],
            c="purple",
            label=f"Agent {agent_j.id} Path",
        )
        plt.scatter(
            agent_j.start[0],
            agent_j.start[1],
            c="r",
            marker="o",
            label=f"Agent {agent_j.id} Start",
        )
        plt.scatter(
            agent_j.goal[0],
            agent_j.goal[1],
            c="r",
            marker="x",
            label=f"Agent {agent_j.id} Goal",
        )
        plt.scatter(
            start_conflict_j[0],
            start_conflict_j[1],
            c="r",
            marker="s",
            label=f"Agent {agent_j.id} Conflict Point",
        )
        plt.scatter(
            end_conflict_j[0],
            end_conflict_j[1],
            c="r",
            marker="p",
            label=f"Agent {agent_j.id} End Conflict",
        )

        # 添加图例和标签
        plt.legend()
        plt.xlabel("X Position")
        plt.ylabel("Y Position")
        plt.title(
            f"Conflict between Agent {agent_i.id} and Agent {agent_j.id} at step {idx}"
        )
        plt.grid(True)

        # 显示绘图
        plt.show()

    def safe_distance(
        self, solution: Dict[Agent, np.ndarray], agent_i: Agent, agent_j: Agent
    ) -> int:
        # 加速 0 2 减速 1 3
        point_start_if_need = -1
        parallel_start = -1
        for idx, (point_i, point_j) in enumerate(
            zip(solution[agent_i], solution[agent_j])
        ):
            distance = self.dist(point_i, point_j)

            if point_start_if_need == -1:
                point_start_if_need = idx if distance <= 2.5 * self.robot_radius else -1
            if parallel_start == -1:
                parallel_start = idx if distance <= 5 * self.robot_radius else -1
            if distance <= 2 * self.robot_radius:
                print(
                    f"智能体{agent_i.id}和智能体{agent_j.id}在第{idx}个位置发生冲突，距离为{self.dist(point_i, point_j)}  "
                )
                print(
                    f"智能体{agent_i.id}起始位置为{agent_i.start}，目标位置为{agent_i.goal}"
                )
                print(
                    f"智能体{agent_j.id}起始位置为{agent_j.start}，目标位置为{agent_j.goal}"
                )
                print(
                    f"发生冲突位置，智能体{agent_i.id}位置为{point_i}，智能体{agent_j.id}位置为{point_j}"
                )

                contrained_path = solution[agent_i]
                unchanged_path = solution[agent_j]

                pivot = unchanged_path[idx]
                conflict_end_time = idx

                try:
                    while (
                        self.dist(contrained_path[conflict_end_time], pivot)
                        <= 2 * self.robot_radius
                    ):
                        conflict_end_time += 1
                except IndexError:
                    pass
                # self.plot_conflict(
                #     solution[agent_i],
                #     solution[agent_j],
                #     agent_i,
                #     agent_j,
                #     point_start_if_need,
                #     conflict_end_time,
                # )
                if point_i[0] < point_j[0]:
                    conflict_slowdown = (
                        1,
                        contrained_path[point_start_if_need][0],
                        contrained_path[point_start_if_need][1],
                        contrained_path[conflict_end_time][0],
                        contrained_path[conflict_end_time][1],
                    )
                    conflict_accelerate = (
                        0,
                        unchanged_path[point_start_if_need][0],
                        unchanged_path[point_start_if_need][1],
                        unchanged_path[conflict_end_time][0],
                        unchanged_path[conflict_end_time][1],
                    )
                    return idx, [
                        "point",
                        "i_j",
                        conflict_accelerate,
                        conflict_slowdown,
                    ]
                else:
                    conflict_accelerate = (
                        0,
                        contrained_path[point_start_if_need][0],
                        contrained_path[point_start_if_need][1],
                        contrained_path[conflict_end_time][0],
                        contrained_path[conflict_end_time][1],
                    )
                    conflict_slowdown = (
                        1,
                        unchanged_path[point_start_if_need][0],
                        unchanged_path[point_start_if_need][1],
                        unchanged_path[conflict_end_time][0],
                        unchanged_path[conflict_end_time][1],
                    )
                    return idx, [
                        "point",
                        "j_i",
                        conflict_accelerate,
                        conflict_slowdown,
                    ]

            # 到此处，遇到了距离小于2*robot_radius的点
            # 判断是不是当前位置前往纵坐标只需要向右移动，冲突是因为agent_i在agent_j的左侧，而目标点的agent_i在agent_j的右侧
            elif (
                abs(point_i[1] - point_j[1]) < 10
                and distance <= 5 * self.robot_radius
                and agent_i.goal[1] == agent_j.goal[1]
            ):
                print(
                    f"智能体{agent_i.id}和智能体{agent_j.id}在第{idx}个位置目标交叉冲突，距离为{self.dist(point_i, point_j)} "
                )
                print(
                    f"智能体{agent_i.id}起始位置为{agent_i.start}，目标位置为{agent_i.goal}"
                )
                print(
                    f"智能体{agent_j.id}起始位置为{agent_j.start}，目标位置为{agent_j.goal}"
                )
                print(
                    f"发生冲突位置，智能体{agent_i.id}位置为{point_i}，智能体{agent_j.id}位置为{point_j}"
                )

                contrained_path = solution[agent_i]
                unchanged_path = solution[agent_j]

                pivot = unchanged_path[idx]
                conflict_end_time = idx

                try:
                    while (
                        self.dist(contrained_path[conflict_end_time], pivot)
                        <= 2.5 * self.robot_radius
                        or abs(contrained_path[conflict_end_time][0] - pivot[0]) < 100
                    ):
                        conflict_end_time += 1
                except IndexError:
                    pass
                self.plot_conflict(
                    solution[agent_i],
                    solution[agent_j],
                    agent_i,
                    agent_j,
                    parallel_start,
                    conflict_end_time,
                )
                if point_i[0] < point_j[0] and agent_i.goal[0] > agent_j.goal[0]:
                    conflict_slowdown = (
                        2,
                        contrained_path[parallel_start][0],
                        contrained_path[parallel_start][1],
                        contrained_path[conflict_end_time][0],
                        contrained_path[conflict_end_time][1],
                    )
                    conflict_accelerate = (
                        3,
                        unchanged_path[parallel_start][0],
                        unchanged_path[parallel_start][1],
                        unchanged_path[conflict_end_time][0],
                        unchanged_path[conflict_end_time][1],
                    )

                    return idx, [
                        "parallel",
                        "i_j",
                        conflict_accelerate,
                        conflict_slowdown,
                    ]
                else:
                    conflict_accelerate = (
                        3,
                        contrained_path[point_start_if_need][0],
                        contrained_path[point_start_if_need][1],
                        contrained_path[conflict_end_time][0],
                        contrained_path[conflict_end_time][1],
                    )
                    conflict_slowdown = (
                        2,
                        unchanged_path[point_start_if_need][0],
                        unchanged_path[point_start_if_need][1],
                        unchanged_path[conflict_end_time][0],
                        unchanged_path[conflict_end_time][1],
                    )
                    return idx, [
                        "parallel",
                        "j_i",
                        conflict_accelerate,
                        conflict_slowdown,
                    ]

        return -1, ["", "", -1, -1]

    @staticmethod
    def dist(point1: np.ndarray, point2: np.ndarray) -> int:
        return int(np.linalg.norm(point1 - point2, 2))  # L2 norm

    def calculate_constraints(
        self,
        node: CTNode,
        constrained_agent: Agent,
        unchanged_agent: Agent,
        time_of_conflict: int,
        conflict_space: List[
            Tuple[
                str, str, Tuple[int, int, int, int, int], Tuple[int, int, int, int, int]
            ]
        ],
        caclulate_type: str,
    ) -> Constraints:
        contrained_path = node.solution[constrained_agent]
        unchanged_path = node.solution[unchanged_agent]

        pivot = unchanged_path[time_of_conflict]
        conflict_end_time = time_of_conflict
        try:
            while (
                self.dist(contrained_path[conflict_end_time], pivot)
                <= 2 * self.robot_radius
            ):
                conflict_end_time += 1
        except IndexError:
            pass

        return node.constraints.fork(
            constrained_agent,
            unchanged_agent,
            tuple(pivot.tolist()),
            time_of_conflict,
            conflict_end_time,
            conflict_space,
            caclulate_type,
        )

    def calculate_goal_times(self, node: CTNode, agent: Agent, agents: List[Agent]):
        solution = node.solution
        goal_times = dict()
        for other_agent in agents:
            if other_agent == agent:
                continue
            time = len(solution[other_agent]) - 1
            goal_times.setdefault(time, set()).add(tuple(solution[other_agent][time]))
        return goal_times

    """
    Calculate the paths for all agents with space-time constraints
    """

    def calculate_path(
        self,
        agent: Agent,
        constraints: Constraints,
        goal_times: Dict[int, Set[Tuple[int, int]]],
    ) -> np.ndarray:
        # print("Calculating path for agent", agent, agent.id)
        return self.st_planner.plan(
            agent.start,
            agent.goal,
            constraints.setdefault(agent, dict()),
            constraints.setdefault_speed(agent, set()),
            agent.id,
            semi_dynamic_obstacles=goal_times,
            max_iter=self.low_level_max_iter,
            debug=self.debug,
        )

    """
    Reformat the solution to a numpy array
    """

    @staticmethod
    def reformat(agents: List[Agent], solution: Dict[Agent, np.ndarray]):
        solution = Planner.pad(solution)
        reformatted_solution = {}
        for agent in agents:
            reformatted_solution[agent] = solution[agent]
        return reformatted_solution

    """
    Pad paths to equal length, inefficient but well..
    """

    @staticmethod
    def pad(solution: Dict[Agent, np.ndarray]):
        max_ = max(max(len(path) for path in solution.values()), 100)
        for agent, path in solution.items():
            if len(path) == max_:
                continue
            padded = np.concatenate(
                [path, np.array(list([path[-1]]) * (max_ - len(path)))]
            )
            solution[agent] = padded
        return solution
